import os
import pandas as pd
import re
import json
from config import DATA_DIR
from tqdm import tqdm

def string_clean(value):
    """统一清理字符串：去除空白，处理特殊占位符"""  
    value = re.sub(r'\s+', '', value.strip())
    if value in ["未知", "", "-", "_"]:
        return None  
    return value

def age_clean(age):
    """处理年龄字段：转为整数或None"""
    if age is None:
        return None
    
    try:
        return int(age)
    except (ValueError, TypeError):
        return None
    
def birthplace_clean(value):
    if value and isinstance(value, str):
        value = value.replace("中国-", "").replace("省", "").replace("市", "")
    return value

def schools_clean(school_str):
    """处理毕业院校：分割字符串并清理"""
    if not school_str:
        return None
    
    # 按分隔符切割
    schools = re.split(r'[、+]', school_str)
    
    # 清理每个学校名称
    cleaned_schools = []
    for school in schools:
        if school:
            # 特殊处理：上海交大 -> 上交大
            if school == "上海交大":
                cleaned = "上交大"
            else:
                # 普通处理：去掉"大学"字样（如果存在）
                cleaned = school.replace("大学", "")
            cleaned_schools.append(cleaned)
    
    return cleaned_schools if cleaned_schools else None

def industry_clean(industry_str):
    """处理所属行业：分割字符串"""
    if not industry_str:
        return None
    
    industries = re.split(r'[、+]', industry_str)
    return [i for i in industries if i] or None

def wealth_clean(value):
    try:
        value = float(value) if value not in [None, ""] else None
    except (ValueError, TypeError):
        value = None  
    return value

def ranking_change_clean(value):
    """处理排名变化字段：转为整数或None"""
    if isinstance(value, str) and value.upper() == "NEW":
        return None 
    try:
        return int(value)
    except (ValueError, TypeError):
        return None

def wealth_change_clean(value):
    """处理财富变化百分比：转为浮点数或None"""
    if isinstance(value, str) and value.upper() == "NEW":
        return None 
    try:
        return float(value.replace('%', '')) / 100 
    except (ValueError, TypeError):
        return None


def data_cleaning(entry):
    """清洗每条数据"""

    #  提取原始字段
    character_info = entry["hs_Character"][0]
    raw_info = {
        "gender": character_info.get("hs_Character_Gender", ""),
        "age": character_info.get("hs_Character_Age", ""),
        "birthplace": character_info.get("hs_Character_BirthPlace_Cn", "").replace("中国-", ""),
        "education": character_info.get("hs_Character_Education_Cn", ""),
        "school": character_info.get("hs_Character_School_Cn", ""),
        "wealth": entry.get("hs_Rank_Rich_Wealth", 0),
        "industry": entry.get("hs_Rank_Rich_Industry_Cn", ""),
        "ranking_change":entry.get("hs_Rank_Rich_Ranking_Change", ""),
        "wealth_change":entry.get("hs_Rank_Rich_Wealth_Change", "")
    }

    # 数据清洗
    cleaned_info={}
    for key, value in raw_info.items():
        if isinstance(value, str):
            value = string_clean(value)
        if key == "birthplace":
            value = birthplace_clean(value)
        elif key == "age":
            value = age_clean(value)
        elif key == "school":
            value = schools_clean(value)
        elif key == "industry":
            value = industry_clean(value)
        elif key == "wealth":
            value = wealth_clean(value)
        elif key == "ranking_change":
            value = ranking_change_clean(value)
        elif key == "wealth_change":
            value = wealth_change_clean(value)  
        cleaned_info[key]=value
  
    return cleaned_info


raw_data_path=os.path.join(DATA_DIR,'raw','hurun_rich_list_2024.json')
with open(raw_data_path, 'r', encoding='utf-8') as f:
    raw_data = json.load(f)  

print("开始清洗数据......")

cleaned_data = [data_cleaning(entry) for entry in raw_data]

print("数据清洗完成！")

save_path=os.path.join(DATA_DIR,'processed','hurun_rich_list_cleaned.csv')
df = pd.DataFrame(cleaned_data)
df.to_csv(save_path, index=False, encoding='utf-8-sig')
print(f"文件已保存至：{save_path}")